{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 准备数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers, optimizers, datasets\n",
    "\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # or any {'0', '1', '2'}\n",
    "\n",
    "def mnist_dataset():\n",
    "    (x, y), (x_test, y_test) = datasets.mnist.load_data()\n",
    "    #normalize\n",
    "    x = x/255.0\n",
    "    x_test = x_test/255.0\n",
    "    \n",
    "    return (x, y), (x_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(1, 'a'), (2, 'b'), (3, 'c'), (4, 'd')]\n"
     ]
    }
   ],
   "source": [
    "print(list(zip([1, 2, 3, 4], ['a', 'b', 'c', 'd'])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 建立模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class myModel:\n",
    "    def __init__(self):\n",
    "        ####################\n",
    "        self.W1 = tf.Variable(tf.initializers.he_normal()([784, 512]), name='W1')\n",
    "        self.b1 = tf.Variable(tf.zeros([512]), name='b1')\n",
    "        self.W2 = tf.Variable(tf.initializers.he_normal()([512, 256]), name='W2')\n",
    "        self.b2 = tf.Variable(tf.zeros([256]), name='b2')\n",
    "        self.W3 = tf.Variable(tf.initializers.he_normal()([256, 10]), name='W3')\n",
    "        self.b3 = tf.Variable(tf.zeros([10]), name='b3')\n",
    "        ####################\n",
    "    def __call__(self, x):\n",
    "        ####################\n",
    "        x = tf.reshape(x, [-1, 784])\n",
    "        h = tf.nn.relu(tf.matmul(x, self.W1) + self.b1)\n",
    "        h = tf.nn.relu(tf.matmul(h, self.W2) + self.b2)\n",
    "        logits = tf.matmul(h, self.W3) + self.b3\n",
    "        ####################\n",
    "        return logits\n",
    "        \n",
    "model = myModel()\n",
    "\n",
    "optimizer = optimizers.Adam()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 计算 loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def compute_loss(logits, labels):\n",
    "    return tf.reduce_mean(\n",
    "        tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "            logits=logits, labels=labels))\n",
    "\n",
    "@tf.function\n",
    "def compute_accuracy(logits, labels):\n",
    "    predictions = tf.argmax(logits, axis=1)\n",
    "    return tf.reduce_mean(tf.cast(tf.equal(predictions, labels), tf.float32))\n",
    "\n",
    "@tf.function\n",
    "def train_one_step(model, optimizer, x, y):\n",
    "    with tf.GradientTape() as tape:\n",
    "        logits = model(x)\n",
    "        loss = compute_loss(logits, y)\n",
    "\n",
    "    # compute gradient\n",
    "    trainable_vars = [model.W1, model.W2, model.b1, model.b2]\n",
    "    grads = tape.gradient(loss, trainable_vars)\n",
    "    for g, v in zip(grads, trainable_vars):\n",
    "        v.assign_sub(0.01*g)\n",
    "\n",
    "    accuracy = compute_accuracy(logits, y)\n",
    "\n",
    "    # loss and accuracy is scalar tensor\n",
    "    return loss, accuracy\n",
    "\n",
    "@tf.function\n",
    "def test(model, x, y):\n",
    "    logits = model(x)\n",
    "    loss = compute_loss(logits, y)\n",
    "    accuracy = compute_accuracy(logits, y)\n",
    "    return loss, accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 实际训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 0 : loss 0.402537 ; accuracy 0.89496666\n",
      "epoch 1 : loss 0.4023711 ; accuracy 0.89498335\n",
      "epoch 2 : loss 0.40220553 ; accuracy 0.895\n",
      "epoch 3 : loss 0.4020403 ; accuracy 0.895\n",
      "epoch 4 : loss 0.40187535 ; accuracy 0.89503336\n",
      "epoch 5 : loss 0.40171075 ; accuracy 0.89505\n",
      "epoch 6 : loss 0.40154648 ; accuracy 0.8950667\n",
      "epoch 7 : loss 0.40138245 ; accuracy 0.8950833\n",
      "epoch 8 : loss 0.40121874 ; accuracy 0.8951333\n",
      "epoch 9 : loss 0.4010554 ; accuracy 0.8950667\n",
      "epoch 10 : loss 0.40089226 ; accuracy 0.8951\n",
      "epoch 11 : loss 0.40072948 ; accuracy 0.8951167\n",
      "epoch 12 : loss 0.400567 ; accuracy 0.89515\n",
      "epoch 13 : loss 0.40040484 ; accuracy 0.8951667\n",
      "epoch 14 : loss 0.40024295 ; accuracy 0.8951833\n",
      "epoch 15 : loss 0.40008146 ; accuracy 0.8952\n",
      "epoch 16 : loss 0.39992023 ; accuracy 0.89525\n",
      "epoch 17 : loss 0.39975926 ; accuracy 0.89528334\n",
      "epoch 18 : loss 0.39959866 ; accuracy 0.89528334\n",
      "epoch 19 : loss 0.39943832 ; accuracy 0.89531666\n",
      "epoch 20 : loss 0.39927828 ; accuracy 0.89535\n",
      "epoch 21 : loss 0.39911854 ; accuracy 0.8954\n",
      "epoch 22 : loss 0.3989591 ; accuracy 0.8954667\n",
      "epoch 23 : loss 0.39879996 ; accuracy 0.8955167\n",
      "epoch 24 : loss 0.39864108 ; accuracy 0.89555\n",
      "epoch 25 : loss 0.39848253 ; accuracy 0.89556664\n",
      "epoch 26 : loss 0.39832425 ; accuracy 0.89558333\n",
      "epoch 27 : loss 0.39816627 ; accuracy 0.89561665\n",
      "epoch 28 : loss 0.3980086 ; accuracy 0.8957\n",
      "epoch 29 : loss 0.39785123 ; accuracy 0.8957667\n",
      "epoch 30 : loss 0.3976942 ; accuracy 0.8958167\n",
      "epoch 31 : loss 0.39753744 ; accuracy 0.8958\n",
      "epoch 32 : loss 0.39738095 ; accuracy 0.8958333\n",
      "epoch 33 : loss 0.39722478 ; accuracy 0.8958833\n",
      "epoch 34 : loss 0.3970689 ; accuracy 0.8958833\n",
      "epoch 35 : loss 0.3969133 ; accuracy 0.8959\n",
      "epoch 36 : loss 0.39675793 ; accuracy 0.89591664\n",
      "epoch 37 : loss 0.39660287 ; accuracy 0.896\n",
      "epoch 38 : loss 0.3964481 ; accuracy 0.896\n",
      "epoch 39 : loss 0.3962936 ; accuracy 0.89598334\n",
      "epoch 40 : loss 0.39613938 ; accuracy 0.89593333\n",
      "epoch 41 : loss 0.39598545 ; accuracy 0.89595\n",
      "epoch 42 : loss 0.39583176 ; accuracy 0.89593333\n",
      "epoch 43 : loss 0.39567843 ; accuracy 0.89596665\n",
      "epoch 44 : loss 0.39552534 ; accuracy 0.89598334\n",
      "epoch 45 : loss 0.39537257 ; accuracy 0.89601666\n",
      "epoch 46 : loss 0.39522004 ; accuracy 0.89603335\n",
      "epoch 47 : loss 0.3950678 ; accuracy 0.89608335\n",
      "epoch 48 : loss 0.39491582 ; accuracy 0.89608335\n",
      "epoch 49 : loss 0.3947641 ; accuracy 0.8961\n",
      "test loss 0.3759438 ; accuracy 0.9017\n"
     ]
    }
   ],
   "source": [
    "train_data, test_data = mnist_dataset()\n",
    "for epoch in range(50):\n",
    "    loss, accuracy = train_one_step(model, optimizer, \n",
    "                                    tf.constant(train_data[0], dtype=tf.float32), \n",
    "                                    tf.constant(train_data[1], dtype=tf.int64))\n",
    "    print('epoch', epoch, ': loss', loss.numpy(), '; accuracy', accuracy.numpy())\n",
    "loss, accuracy = test(model, \n",
    "                      tf.constant(test_data[0], dtype=tf.float32), \n",
    "                      tf.constant(test_data[1], dtype=tf.int64))\n",
    "\n",
    "print('test loss', loss.numpy(), '; accuracy', accuracy.numpy())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
